package com.shouke.jmeter.plugin.sampler;

import com.shouke.RabbitMQ.RabbitMQClient;
import com.shouke.jmeter.plugin.model.FieldExpressionMapping;
import org.apache.jmeter.samplers.Entry;
import org.apache.jmeter.samplers.SampleResult;
import org.apache.jmeter.samplers.Sampler;
import org.apache.jmeter.testbeans.TestBean;
import org.apache.jmeter.testelement.AbstractTestElement;
import org.apache.jmeter.testelement.TestStateListener;
import org.apache.jmeter.threads.JMeterContextService;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.json.JSONObject;

import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


public class RabbitMQPubSampler extends AbstractTestElement implements Sampler, TestBean, TestStateListener {
    private static final Logger logger = LogManager.getLogger(RabbitMQPubSampler.class);
    private static final long serialVersionUID = 16886449268470L;

    private String rabbitMQConnVariableName;

    // 队列设置
    private String exchangeName; // 交换机名称
    private String queueName; // 队列名称
    private Boolean isCreateQueue; // 是否创建队列
    private String routingKey; // routing key
    private List<FieldExpressionMapping> messageHeaders; // 消息头
    private List<FieldExpressionMapping> messageProperties; // 消息属性
    private String dataType; // 数据类型
    private String message; // 消息

    private RabbitMQClient rabbitMQClient;

    @Override
    public SampleResult sample(Entry e) {
        SampleResult result = new SampleResult();
        try {
            if(this.rabbitMQClient == null){
                this.rabbitMQClient = getRabbitMQClient();
            }
            Object[] tempRes = handleMessageHanders();
            String headersInStr = (String)tempRes[0];
            Map<String, Object> headers = (Map<String, Object>)tempRes[1];

            tempRes = handleMessageProperties();
            String propertiesInStr = (String)tempRes[0];
            Map<String, String> properties = (Map<String, String>)tempRes[1];

            String routingKey = this.routingKey;
            if (this.routingKey.isEmpty() && !this.queueName.isEmpty()) {
                routingKey = this.queueName;
            }
            if (!propertiesInStr.contains("contentType") && this.dataType.equals("JSON")) {
                propertiesInStr = propertiesInStr + "contentType: application/json";
            }

            String SampleData = "Exchange: " + this.exchangeName + "\nRoutingKey: " + routingKey  + "\nProperties: \n" + propertiesInStr + "\nMessage: \n" + getMessage();

            result.setSampleLabel(getName());
            result.setDataType(SampleResult.TEXT);
            result.setContentType("");
            result.setDataEncoding(StandardCharsets.UTF_8.name());
            result.setSamplerData(SampleData);
            result.setRequestHeaders(headersInStr);
            result.sampleStart();
            String res = this.rabbitMQClient.basicPublish(this.exchangeName, this.queueName, this.isCreateQueue, routingKey, headers, properties, this.dataType, this.message);
            result.setResponseData(res, StandardCharsets.UTF_8.name());
            result.setResponseHeaders("");
            result.setResponseOK();
        } catch (Exception ex) {
            logger.error("Exception occurred while executing publishing message");
            result = handleException(result, ex);
        } finally {
            result.sampleEnd();
        }
        return result;
    }

    /**
     * 主线程开始执行测试之前被调用，注意,不管有多少线程，循环运行多少次，该函数仅会被调用一次
     * 优先于配置元件中的testStarted()函数被执行
     */
    @Override
    public void testStarted() {
    }

    @Override
    public void testStarted(String host) {
        logger.debug("sampler testStarted, arg：" + host);
        this.testStarted();
    }

    /**
     * 测试结束后被调用，注意，不管有多少线程，循环运行多少次，该函数仅会被调用一次
     * 优先于配置元件中的testEnded()函数被执行
     */
    @Override
    public void testEnded() {
    }

    @Override
    public void testEnded(String host) {
        logger.debug("sampler testEnded, arg：" + host);
    }

    public String getRabbitMQConnVariableName() {
        return rabbitMQConnVariableName;
    }

    public void setRabbitMQConnVariableName(String rabbitMQConnVariableName) {
        this.rabbitMQConnVariableName = rabbitMQConnVariableName;
    }

    private RabbitMQClient getRabbitMQClient() {
        return (RabbitMQClient) JMeterContextService.getContext().getVariables().getObject(getRabbitMQConnVariableName());
    }


    private SampleResult handleException(SampleResult result, Exception ex) {
        result.setResponseMessage("Error occured while publishing message");
        result.setResponseCode("500");
        result.setResponseData(String.format("Error occured while publishing message：%s", ex.toString()).getBytes());
        result.setSuccessful(false);
        return result;
    }

    public void setExchangeName(String exchangeName) {
        this.exchangeName = exchangeName.trim();
    }

    public String getExchangeName(){
        return this.exchangeName;
    }

    public String getQueueName(){
        return this.queueName;
    }

    public void setQueueName(String queueName) {
        this.queueName = queueName.trim();
    }

    public Boolean getIsCreateQueue(){
        return this.isCreateQueue;
    }

    public void setIsCreateQueue(Boolean isCreateQueue) {
        this.isCreateQueue = isCreateQueue;
    }

    public String getRoutingKey(){
        return this.routingKey;
    }

    public void setRoutingKey(String routingKey) {
        this.routingKey = routingKey.trim();
    }

    public List<FieldExpressionMapping> getMessageHeaders(){
        return this.messageHeaders;
    }

    public void setMessageHeaders(List<FieldExpressionMapping> messageHeaders) {
        this.messageHeaders = messageHeaders;
    }

    public List<FieldExpressionMapping> getMessageProperties(){
        return this.messageProperties;
    }

    public void setMessageProperties(List<FieldExpressionMapping> messageProperties) {
        this.messageProperties = messageProperties;
    }

    public String getDataType(){
        return this.dataType;
    }

    public void setDataType(String dataType) {
        this.dataType = dataType.trim();
    }

    public String getMessage(){
        return this.message;
    }

    public void setMessage(String message) {
        this.message = message.trim();
        if (this.dataType.equals("JSON")) {
            JSONObject jsonObject = new JSONObject(this.message);
            this.message = jsonObject.toString(2);
        }
    }

    public Object[] handleMessageHanders() {
        // The map for the headers.
        Map<String, Object> headers = new HashMap<String, Object>();
        StringBuffer sb = new StringBuffer();

        for (FieldExpressionMapping entry : getMessageHeaders()) {
            String header = entry.getMsgHeader().trim();
            String value = entry.getMsgHeaderValue().trim();
            headers.put(header, value);
            sb.append(header).append(":").append(value).append("\n");
        }
        return new Object[]{ sb.toString(), headers};
    }

    public Object[] handleMessageProperties() {
        Map<String, String> properties = new HashMap<String, String>();
        StringBuffer sb = new StringBuffer();

        for (FieldExpressionMapping entry : getMessageProperties()) {
            String property = entry.getMsgProperty().trim();
            String value = entry.getMsgPropertyValue().trim();
            if (property.equals("contentType") && this.dataType.equals("JSON")) {
                value = "application/json";
            }
            properties.put(property, value);
            sb.append(property).append(":").append(value).append("\n");
        }
        return new Object[]{ sb.toString(), properties};    }

}